# -*- coding: utf-8 -*-
"""
@Author：mengying
@file： worker.py
@date：2023/7/24 9:27
@email: 652044581@qq.com
@desc: 
"""
from redis import Redis
from playwright.sync_api import sync_playwright, Playwright
from threading import Thread
import logging
import os
import base64
import platform
import time
import json
from typing import Union
from extract import JD_extract

logger = logging.getLogger(__file__)


class config:
    # 配置redis缓存地址
    REDIS_HOST = '124.71.212.219'
    REDIS_DB = 1
    REDIS_PORT = 6379
    REDIS_PASSWORD = 'YFKUDeCWD1AL3pGnxHtVoAf3dcwcrTgL'
    WORKER_QUEUE = "WORKER_QUEUE_PLAYWRIGHT"


# 配置缓存
RedisClient: Redis = Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, decode_responses=True,
                           password=config.REDIS_PASSWORD, db=config.REDIS_DB)

BASE_PATH = os.path.join(os.path.dirname(__file__), "html")


class Formatter:

    @staticmethod
    def tobase64(value: Union[str, bytes]):
        content = value if isinstance(value, bytes) else value.encode()
        return base64.b64encode(content).decode()


class DriverWorker:

    def __init__(self, playwright: Playwright):
        headless = True
        if 'windows' in platform.platform().lower():
            headless = False
        self.browser = playwright.chromium.launch(headless=headless)
        self.context = self.browser.new_context()
        self.page = self.context.new_page()

    def run(self, url):
        self.page.goto(url)

    def get_content(self):
        return self.page.content()

    def get_pdf(self):
        return self.page.pdf()


class ManagerWorker:

    @staticmethod
    def run():
        while True:
            try:
                worker_lens = RedisClient.llen(config.WORKER_QUEUE)

                # 判断有无任务队列
                if not worker_lens:
                    time.sleep(1)
                    continue

                # 启动任务
                with sync_playwright() as playwright:
                    driver = DriverWorker(playwright)

                    while True:
                        worker_lens = RedisClient.llen(config.WORKER_QUEUE)
                        if worker_lens:
                            for i in range(worker_lens):
                                try:
                                    data = RedisClient.brpop(config.WORKER_QUEUE, timeout=30)
                                except Exception as e:
                                    driver.browser.close()
                                else:
                                    data = json.loads(data[1])
                                    url = data.get('url')
                                    content = data.get('content')
                                    mode = data.get('mode')
                                    uid = data.get('uid')
                                    ret = ""
                                    logger.info("任务worker: %s 开始" % uid)
                                    time1 = time.time()
                                    if url:
                                        driver.run(url)
                                        if str(mode) == "0":
                                            ret = driver.get_content()
                                            ret = str(JD_extract.get_web_data(ret))
                                        elif str(mode) == "1":
                                            ret = driver.get_pdf()
                                    else:
                                        html = base64.b64decode(content).decode("utf-8")
                                        save_path = os.path.join(BASE_PATH, uid + ".html")
                                        with open(save_path, "w") as f:
                                            f.write(html)
                                        driver.run(save_path)
                                        if str(mode) == "0":
                                            ret = driver.get_content()
                                        elif str(mode) == "1":
                                            ret = driver.get_pdf()

                                    # 存放数据
                                    ret = Formatter.tobase64(ret)
                                    RedisClient.set(uid, ret, ex=60 * 10)
                                    time2 = time.time()
                                    logger.info("任务worker: %s 结束 耗时：%s" % (uid, str(time2 - time1)))
                        else:
                            driver.browser.close()
            except Exception as e:
                continue


Thread(target=ManagerWorker.run).start()
